Skip to content

Update trtllm-gen fused moe routing kernel and add more kernels #1955

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Draft
wants to merge 12 commits into
base: main
Choose a base branch
from
Draft

Update trtllm-gen fused moe routing kernel and add more kernels #1955

wants to merge 12 commits into from

Conversation

@jiahanc
Copy link
Collaborator

@jiahanc jiahanc commented Oct 20, 2025 •
edited by IwakuraRein
Loading

📌 Description

  • update the trtllm-gen fused moe headers
  • add new kernels for trtllm-gen fused moe
    • for NvFp4, add tile 256
    • for MxFp8 x MxFp4, add 128, 256
    • for FP8 per-tensor, add 192, 256
    • for FP8 block scale, add 128
  • update the logics of computeSelectedTileN
  • add tune_max_num_tokens to FP8 per-tensor and FP8 block scale

🔍 Related Issues

🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.

✅ Pre-commit Checks

  • I have installed pre-commit by running pip install pre-commit (or used your preferred method).
  • I have installed the hooks with pre-commit install.
  • I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

If you are unsure about how to set up pre-commit, see the pre-commit documentation.

🧪 Tests

  • Tests have been added or updated as needed.
  • All tests are passing (unittest, etc.).

Reviewer Notes

@coderabbitai
Copy link
Contributor

coderabbitai bot commented Oct 20, 2025 •
edited
Loading

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

✨ Finishing touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Post copyable unit tests in a comment

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

@jiahanc jiahanc force-pushed the updateTileNCalc branch 2 times, most recently from 9d9ad95 to 7cd156d Compare October 22, 2025 22:21
@jiahanc jiahanc mentioned this pull request Oct 28, 2025
Merged
5 tasks
@IwakuraRein
update
c833388 
Signed-off-by: Siyuan Fu <[email protected]>
@IwakuraRein
update
d20582f 
Signed-off-by: Siyuan Fu <[email protected]>
@IwakuraRein
bump trtllm-gen to ac83a
78004cf 
Signed-off-by: Siyuan Fu <[email protected]>
@IwakuraRein
public cubins
8935688
@IwakuraRein IwakuraRein changed the title Revise the calculation related to TileN in routing of MOE TRTLLM backend Update trtllm-gen fused moe routing kernel and add more kernels Nov 1, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@joker-eph joker-eph Awaiting requested review from joker-eph joker-eph will be requested when the pull request is marked ready for review joker-eph is a code owner

@aleozlx aleozlx Awaiting requested review from aleozlx aleozlx will be requested when the pull request is marked ready for review aleozlx is a code owner

@yzh119 yzh119 Awaiting requested review from yzh119 yzh119 will be requested when the pull request is marked ready for review yzh119 is a code owner

@cyx-6 cyx-6 Awaiting requested review from cyx-6 cyx-6 will be requested when the pull request is marked ready for review cyx-6 is a code owner

@wenscarl wenscarl Awaiting requested review from wenscarl wenscarl will be requested when the pull request is marked ready for review wenscarl is a code owner

@djmmoss djmmoss Awaiting requested review from djmmoss djmmoss will be requested when the pull request is marked ready for review djmmoss is a code owner

@yongwww yongwww Awaiting requested review from yongwww yongwww will be requested when the pull request is marked ready for review yongwww is a code owner

@IwakuraRein IwakuraRein Awaiting requested review from IwakuraRein IwakuraRein will be requested when the pull request is marked ready for review IwakuraRein is a code owner

@nvmbreughe nvmbreughe Awaiting requested review from nvmbreughe nvmbreughe will be requested when the pull request is marked ready for review nvmbreughe is a code owner

@bkryu bkryu Awaiting requested review from bkryu bkryu will be requested when the pull request is marked ready for review bkryu is a code owner

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@jiahanc @IwakuraRein